#coding=utf-8
'''
Created on 2011-3-14

@author: zarra
读取mit ECG数据库
'''
import os
from pylab import *
import numpy as np

class EcgChannel(object):
    '''
           实际的ECG数据
    '''
    def __init__(self,data,**kwag):
        name      = kwag.get('name','')
        zerovalue = kwag.get('zerovalue',1024)
        gain      = kwag.get('gain',200)
        self.real_data = np.array(data)
        self.real_data = self.real_data /gain - zerovalue/gain
        self.name = name
    def __getitem__(self,index):
        return self.real_data[index]
    def __getslice__(self,i,j):
        return self.real_data[i:j]
    def __repr__(self):
        return self.real_data.__repr__()
    
class EcgHead(object):
    '''
    Head 文件信息
    '''
    def __init__(self,filename):
        with open(filename,'r') as head:
                line1 = head.readline()
                fields = line1.split(' ')
                self.id = fields[0]
                self.channelCount = int(fields[1])#通道数
                self.sampleFreq = int(fields[2])#采样频率
                self.channels =[]#通道信息
                for i in range(self.channelCount):                 
                    line = head.readline()
                    self.channels.append(self.readChannelConfigLine(line))
                #line3 = head.readline()
        
    def readChannelConfigLine(self,line):
        fields = line.split(' ')
        return {
                'name':fields[8], #导联名称
                'fileName':fields[0],#数据文件名
                'format':fields[1],#数据格式 必须是'212'
                'gain':int(fields[2]),#每 mV 包含的整数个数
                'bitres':int(fields[3]),#采样精度（位分辨率）
                'zerovalue':int(fields[4]),# ECG 信号零点相应的整数值
                'firstvalue':int(fields[5])#信号的第一个整数值 (用于偏差测试)
                }
        
        
class EcgDataFile(object):
    def __init__(self,filename):
        self.head_filename = filename
        self.data_filename = os.path.splitext(filename)[0]+'.dat'
        self.isRead = False
        self.data = []
        if filename != None:
            self.readAll()
    def readAll(self):
        if self.isRead == False:
            data = [[],[]]
            self.head = EcgHead(self.head_filename)
            with open(self.data_filename,'rb') as f:
                while True:
                    array3 = f.read(3)
                    if len(array3) == 3:
                        #每三byte包含了2个通道的各一个数据
                        #这三个byte记作b1,b2,b3
                        #第一byte和第二byte的低4位为通道1的数据
                        #其值为b2[0:3]<<8+b1
                        #第二个byte的高4位和第三byte组成通道二的数据
                        #值为b2[4:7]<<8+b3
                        s1=ord(array3[0])
                        s2=ord(array3[2])
                        m=ord(array3[1])
                        h1= m & 15
                        h2= m>>4
                        r1 = ((h1<<8)+s1)&2047#取低11位
                        r2 = ((h2<<8)+s2)&2047
                        data[0].append(float(r1))
                        data[1].append(float(r2))
                    else:
                        #pass
                        break
            self.isRead = True
            self.data = data
            self.channel=[]
            for i in range(self.head.channelCount):  
                self.channel.append( EcgChannel(self.data[i],**self.head.channels[i]) ) 
        return  self.channel
    def show(self,begin,end,channelID=0):
        channel = self.readAll()[channelID]
        sfreq = self.head.sampleFreq
        x=[float(i)/sfreq for i in range(begin,end)]

        title(channel.name)
        plot(x,channel[begin:end])
        

                
                
if __name__ == '__main__':
    ecgData = EcgDataFile('../data/100.hea')
    print ecgData.head.sampleFreq
    print ecgData.head.channels[1]['firstvalue']
    print ecgData.data[1][0]
    print ecgData.readAll()[0][:10]        
